#lang racket

(require "util/utils.scm")
(require "util/gensym.scm")
(require "util/numberizer.scm")
(require "util/match-assoc.scm")
(provide pass-flat)

;; scheme --script pass-flat.scm t/t1.scm

;; the FLAT language is first order scheme with:
;; - function applications denested
;; - if expressions denested
;; so most of the bodies are LET*s
;;
;; procedure labels exist in a different namespace
;; to top level definition names.
;; every procedure has args and an environment it can
;; reference.
;;
;; DEFINITIONS SECTION
;; (define <name> <label>) *
;;
;; PROCEDURES SECTION
;; (define-procedure (<label> <arg> ...) <num-tmps> <flat>)
;;
;; SYNTAX
;; <flat> ::= (let* ((<tmp> <simple>) ...) <simple>)
;; 
;; <simple> ::= <simpler>
;;            | (apply <simpler> ...)
;;            | (builtin <name> <simpler> ...)
;;            | (if <simpler> <flat> <flat>)
;; 
;; <simpler> ::= (variable <vtype> <var>)
;;             | (quote <datum>)
;;             | (closure (<env> ...) <label>)
;;
;; <vtype> ::= loc | glo | tmp | env

(define (env-lookup env nm i)
  (if (equal? (car env) nm)
      i
      (env-lookup (cdr env) nm (+ i 1))))


(define (make-var-numberizer arg-size tmp-size num-glos env)
  (let ((frame-size (+ arg-size tmp-size))
        (locn (make-numberizer))
        (tmpn (make-numberizer)))
    (define (make locn tmpn)
      (let ((numberize
             (lambda (ty nm stk-offset)
               (case ty
                 ((loc)
                  (let ((n (+ stk-offset (- frame-size 1 (numberize! locn nm)))))
                    (if (< n 0)
                        (error 'numberize-var "negative stack reference" `(variable ,ty ,nm))
                        `(variable stk ,n))))
                 ((tmp)
                  (let ((n (+ stk-offset (- frame-size arg-size 1 (numberize! tmpn nm)))))
                    (if (< n 0)
                        (error 'numberize-var "negative stack reference" `(variable ,ty ,nm)) 
                        `(variable stk ,n))))
                 ((glo) `(variable glo ,(numberize! num-glos nm)))
                 ((env) `(variable env ,(env-lookup env nm 0)))
                 (else (error "unknown variable type" ty))))))
        (lambda xs
          (if (null? xs)
              (make (clone-numberizer locn) (clone-numberizer tmpn))
              (apply numberize xs)))))
    (make locn tmpn)))

(define (pass-flat program)
  (match-assoc program (definitions procedures)
    (let* ((code (make-collector))
           (num-glos (make-numberizer))
	   (definitions (mapply pass-flat-definition definitions code num-glos))
	   (procedures (mapply pass-flat-procedure procedures code num-glos))
	   (start (list (dgensym 'start))))
      (pass-flat-entry-point start definitions code)
      `((entry-point . ,start)
	(globals-size ,(numberizer-length num-glos))
	(code . ,(collector-get code))))))

(define (pass-flat-entry-point start definitions code)
  (collect-object! code `(label ,start))
  (for-each (lambda (def)
	      (match def
		(`(define ,glo (label ,l))
		 (let ((retty (dgensym 'start-cont)))
		   (collect-object! code `(push! (label ,retty)))
		   (collect-object! code `(push! (env-register)))
		   (collect-object! code `(jump ,l))
		   (collect-object! code `(label ,retty))
		   (collect-object! code `(move! (variable glo ,glo) (return-register)))
		   ))))
	    definitions)
  (collect-object! code `(exit)))

(define (pass-flat-definition def code num-glos)
  (match def
    (`(define ,name ,label)
     `(define ,(numberize! num-glos name) ,label))
    (else (error 'pass-flat-definition "?" def))))

(define return-register (lambda (stk-offset) '(return-register)))

(define (pass-flat-procedure def code num-glos)
  (match def
    (`(,label ,env ,args ,num-tmps ,flat)
     (let ((numberize (make-var-numberizer (length args) num-tmps num-glos env))
           (frame-size (+ (length args) num-tmps)))
       ;; label the definition entry point
       ;; and compile the code
       (collect-object! code label)
       ;; numberize the locals first
       (for-each (lambda (a) (numberize 'loc a 0)) args)
       ;; allocate space for the tmps
       (unless (= 0 num-tmps)
	 (collect-object! code `(extend-stack ,num-tmps)))
       ;; compile the code
       (pass-flat-exp flat return-register numberize code 0 frame-size)
       ;;
       (unless (= 0 (+ (length args) num-tmps))
	 (collect-object! code `(shrink-stack ,(+ (length args) num-tmps))))
       
       (collect-object! code `(return))
       ;; make the definition point to that code
       `(define-procedure ,label)))
    (else (error 'pass-flat-procedure "?" def))))

(define (pass-flat-exp exp place numberize code stk-offset frame-size)
  (match exp
    (`(let* ,bindings ,body)
     (mapply pass-flat-into-binding bindings numberize code stk-offset frame-size)
     (pass-flat-into body place numberize code stk-offset frame-size))
    (else (pass-flat-into exp place numberize code stk-offset frame-size))))

(define (pass-flat-into-binding binding numberize code stk-offset frame-size)
  (let ((place (lambda (stk-offset) (numberize 'tmp (car binding) stk-offset))))
    (pass-flat-into (cadr binding) place numberize code stk-offset frame-size)))

(define (pass-flat-into exp place numberize code stk-offset frame-size)
  ;; NOTE: WHAT IF PLACE IS A STK REFERENCE
  ;; need to use stk-offset
  (match exp
    (`(atomic ,x) (collect-object! code `(move! ,(place stk-offset) (atomic ,x))))
    (`(var ,ty ,nm) (collect-object! code `(move! ,(place stk-offset) ,(numberize ty nm stk-offset))))
    (`(quote ,dat) (collect-object! code `(move! ,(place stk-offset) ,exp)))
    (`(closure ,env ,lbl)
     (let ((plce (place stk-offset))
           (len (length env)))
       (collect-object! code `(move! ,plce (closure ,len ,lbl)))
       (let loop ((env env) (n 0))
         (if (null? env)
             '()
             (let* ((e (car env))
                    (val (pass-flat-simpler e numberize stk-offset frame-size)))
               (collect-object! code `(closure-set! ,plce ,n ,val))
               (loop (cdr env) (+ n 1)))))))
    (`(app ,func . ,args)
     (if (eq? place return-register)
         (pass-flat-tail-application #f func args place numberize code stk-offset frame-size)
         (pass-flat-application #f func args place numberize code stk-offset frame-size)))
    (`(known-app ,func . ,args)
     (if (eq? place return-register)
         (pass-flat-tail-application #t func args place numberize code stk-offset frame-size)
         (pass-flat-application #t func args place numberize code stk-offset frame-size)))

    (`(builtin ,func . ,args)
     (collect-object! code `(move! ,(place stk-offset) (builtin ,func . ,(mapply pass-flat-simpler args numberize stk-offset frame-size)))))
    (`(if ,t ,c ,a)
     (let ((lbl-skp (dgensym 'if-skp))
	   (lbl-out (dgensym 'if-out))
           (t1 (pass-flat-simpler t numberize stk-offset frame-size)))
       (collect-object! code `(branch ,t1 ,lbl-skp))
       (pass-flat-exp c place (numberize) code stk-offset frame-size)
       (collect-object! code `(jump ,lbl-out))
       (collect-object! code `(label ,lbl-skp))
       (pass-flat-exp a place (numberize) code stk-offset frame-size)
       (collect-object! code `(label ,lbl-out))))
    (else (error 'pass-flat-into "?" exp))))

(define (pass-flat-tail-application known? func args place numberize code stk-offset frame-size)
  (let ((stk-offset^ stk-offset))
    (for-each (lambda (arg)
                (collect-object! code `(push! ,(pass-flat-simpler arg numberize stk-offset frame-size)))
                (set! stk-offset (+ 1 stk-offset)))
              args)
    (let ((func (if known? func (pass-flat-simpler func numberize stk-offset frame-size)))
          (distance (- (+ frame-size stk-offset) (length args))))
      (if known?
	  (begin
	    (collect-object! code `(stack-fold ,(length args) ,distance))
	    (collect-object! code `(jump ,func)))
	  (begin
	    (collect-object! code `(move! (return-register) ,func))
	    (collect-object! code `(stack-fold ,(length args) ,distance))
	    (collect-object! code `(call))))
      (set! stk-offset stk-offset^))))

(define (pass-flat-application known? func args place numberize code stk-offset frame-size)
  ;; To code an application we have to push the return address
  ;; then we push the args in reverse order
  ;; and leave a continuation label
  (let ((retty (dgensym 'return-address))
        (stk-offset^ stk-offset))
    (collect-object! code `(push! (label ,retty)))
    (set! stk-offset (+ 1 stk-offset))
    (collect-object! code `(push! (env-register)))
    (set! stk-offset (+ 1 stk-offset))
    
    (for-each (lambda (arg)
                (collect-object! code `(push! ,(pass-flat-simpler arg numberize stk-offset frame-size)))
                (set! stk-offset (+ 1 stk-offset)))
              args)
    (if known?
	(collect-object! code `(jump ,func))
	(begin
	  (collect-object! code `(move! (return-register) ,(pass-flat-simpler func numberize stk-offset frame-size)))
	  (collect-object! code `(call))))
    
    (collect-object! code `(label ,retty))
    (set! stk-offset stk-offset^)
    ;; and then we need to move the result of the call into 'place'
    (unless (eq? place return-register)
      (collect-object! code `(move! ,(place stk-offset) (return-register))))))

(define (pass-flat-simpler exp numberize stk-offset frame-size)
  (match exp
    (`(var ,ty ,nm) (numberize ty nm stk-offset))
    (`(quote ,dat) exp)
    (`(atomic ,dat) exp)
    (else (error 'pass-flat-simpler "?" exp))))


